/*  电机编号
            头

     9  5       4 8
        1       0
          |dog|
          |dog|
          |dog|
        3       2
     11 7       6 10
*/

typedef struct point
{
     float x;
     float y;
     float z;
} bezier_Point;

#include "dog_solve.h"

#include <webots/Robot.hpp>
#include <webots/Motor.hpp>
#include <webots/PositionSensor.hpp>
#include <webots/InertialUnit.hpp>
#include <webots/Gyro.hpp>
#include <webots/Accelerometer.hpp>

#include "func.h"
#include <iostream>
#include <string>
#include "all.h"
#include <qpOASES.hpp>
#include <Eigen/Dense>

#include <iomanip>

using namespace std;
using namespace Eigen;
using namespace webots;

// #define NUMBER_OF_LEDS 8
#define NUMBER_OF_JOINTS 12
// #define NUMBER_OF_CAMERAS 5

#define TIME_STEP 100
#define MPC_TEST

static const char *motor_names[NUMBER_OF_JOINTS] = {
    "front right shoulder abduction motor", // 右前肩外展电机
    "front left shoulder abduction motor",  // 左前肩外展电机
    "rear right shoulder abduction motor",  // 右后肩外展电机
    "rear left shoulder abduction motor",   // 左后肩外展电机
    "front right shoulder rotation motor",  // 右前肩旋转电机
    "front left shoulder rotation motor",   // 左前肩旋转电机
    "rear right shoulder rotation motor",   // 右后肩旋转电机
    "rear left shoulder rotation motor",    // 左后肩旋转电机
    "front right elbow motor",              // 右前肘电机
    "front left elbow motor",               // 左前肘电机
    "rear right elbow motor",               // 右后肘电机
    "rear left elbow motor"                 // 左后肘电机
};

static Motor *motors[NUMBER_OF_JOINTS];
static PositionSensor *ps[NUMBER_OF_JOINTS];

void setsmoothly(Robot *robot, const double *target, double duration)
{
     double time_step = robot->getBasicTimeStep();
     const int n_steps_to_achieve_target =
         duration * 1000 / time_step; // 经过多少时间步完成动作

     double step_difference[NUMBER_OF_JOINTS];
     double current_position[NUMBER_OF_JOINTS];
     for (int i = 0; i < NUMBER_OF_JOINTS; ++i)
     {
          current_position[i] = motors[i]->getTargetPosition();
          step_difference[i] = (target[i] - current_position[i]) / n_steps_to_achieve_target;
     }
     for (int i = 0; i < n_steps_to_achieve_target; i++)
     {
          for (int j = 0; j < NUMBER_OF_JOINTS; ++j)
          { // 对每一个电机设置动作
               current_position[j] += step_difference[j];
               motors[j]->setPosition(current_position[j]);
          }
     }
}

// 计算三次贝塞尔曲线上的点
void cubicBezier3D(bezier_Point p0, bezier_Point p1, bezier_Point p2, bezier_Point p3, float t, bezier_Point *result)
{
     result->x = (1 - t) * (1 - t) * (1 - t) * p0.x + 3 * (1 - t) * (1 - t) * t * p1.x + 3 * (1 - t) * t * t * p2.x + t * t * t * p3.x;
     result->y = (1 - t) * (1 - t) * (1 - t) * p0.y + 3 * (1 - t) * (1 - t) * t * p1.y + 3 * (1 - t) * t * t * p2.y + t * t * t * p3.y;
     result->z = (1 - t) * (1 - t) * (1 - t) * p0.z + 3 * (1 - t) * (1 - t) * t * p1.z + 3 * (1 - t) * t * t * p2.z + t * t * t * p3.z;
}

void walk_bezier_forward()
{ // 平地行走

     Robot robot;

     // int timeStep = (int)robot.getBasicTimeStep();
     int timeStep = 100;

     for (int i = 0; i < 12; i++)
     {
          dog_pd[i].kp = 60;
          dog_pd[i].kd = 3;
     }
     // p是起点，控制点1，控制点2，终点
     // 摆动相
     bezier_Point p0 = {-0.200, 0.000, 0.200};
     bezier_Point p1 = {-0.100, 0.000, 0.350};
     bezier_Point p2 = {-0.000, 0.000, 0.350};
     bezier_Point p3 = {0.200, 0.000, 0.200};
     // 另外支撑相
     bezier_Point q0 = {0.200, 0.000, 0.200};
     bezier_Point q1 = {-0.000, 0.000, 0.200};
     bezier_Point q2 = {-0.100, 0.000, 0.200};
     bezier_Point q3 = {-0.200, 0.000, 0.200};

     for (int i = 0; i < NUMBER_OF_JOINTS; i++)
     {
          motors[i] = robot.getMotor(motor_names[i]);
          motors[i]->setPosition(0);
          ps[i] = motors[i]->getPositionSensor();
          ps[i]->enable(TIME_STEP);
     }

     InertialUnit *iu = robot.getInertialUnit("inertial unit");
     iu->enable(TIME_STEP);

     Gyro *gyro = robot.getGyro("gyro");
     gyro->enable(TIME_STEP);

     Accelerometer *acc = robot.getAccelerometer("accelerometer");
     acc->enable(TIME_STEP);

     StateEstimator stateEstimator;
     TrajectoryGenerator trajectoryGenerator;
     GaitPlanner gaitPlanner;
     MPCsolver mpcSolver;
     //   估计本体状态

     while (robot.step(TIME_STEP) != -1)
     {
          for (float t = 0; t <= 1; t += 0.1)
          {
               bezier_Point point_st;
               bezier_Point point_sw;
               cubicBezier3D(p0, p1, p2, p3, t, &point_st); // 14支撑相
               cubicBezier3D(q0, q1, q2, q3, t, &point_sw); // 23摆动相

               foot_position[0].x = point_st.x;
               foot_position[0].y = point_st.y;
               foot_position[0].z = point_st.z;
               foot_position[2].x = point_st.x;
               foot_position[2].y = point_st.y;
               foot_position[2].z = point_st.z;

               foot_position[1].x = point_sw.x;
               foot_position[1].y = point_sw.y;
               foot_position[1].z = point_sw.z;
               foot_position[3].x = point_sw.x;
               foot_position[3].y = point_sw.y;
               foot_position[3].z = point_sw.z;
               dog_solve();

               const double motors_target_pos[NUMBER_OF_JOINTS] = {
                   leg_angle_ouput[0].sol_gamma,
                   leg_angle_ouput[1].sol_gamma,
                   leg_angle_ouput[3].sol_gamma,
                   leg_angle_ouput[2].sol_gamma,
                   leg_angle_ouput[0].sol_alfa,
                   leg_angle_ouput[1].sol_alfa,
                   leg_angle_ouput[3].sol_alfa,
                   leg_angle_ouput[2].sol_alfa,
                   leg_angle_ouput[0].sol_beta,
                   leg_angle_ouput[1].sol_beta,
                   leg_angle_ouput[3].sol_beta,
                   leg_angle_ouput[2].sol_beta};

               setsmoothly(&robot, motors_target_pos, 0.1);
               robot.step(timeStep);
          }

          // 输出100个点
          for (float t = 0; t <= 1; t += 0.1)
          {
               bezier_Point point_sw;
               bezier_Point point_st;
               cubicBezier3D(p0, p1, p2, p3, t, &point_st); // 摆动相
               cubicBezier3D(q0, q1, q2, q3, t, &point_sw); // 支撑相

               foot_position[0].x = point_sw.x;
               foot_position[0].y = point_sw.y;
               foot_position[0].z = point_sw.z;
               foot_position[2].x = point_sw.x;
               foot_position[2].y = point_sw.y;
               foot_position[2].z = point_sw.z;

               foot_position[1].x = point_st.x;
               foot_position[1].y = point_st.y;
               foot_position[1].z = point_st.z;
               foot_position[3].x = point_st.x;
               foot_position[3].y = point_st.y;
               foot_position[3].z = point_st.z;

               dog_solve();

               const double motors_target_pos[NUMBER_OF_JOINTS] = {
                   leg_angle_ouput[0].sol_gamma,
                   leg_angle_ouput[1].sol_gamma,
                   leg_angle_ouput[3].sol_gamma,
                   leg_angle_ouput[2].sol_gamma,
                   leg_angle_ouput[0].sol_alfa,
                   leg_angle_ouput[1].sol_alfa,
                   leg_angle_ouput[3].sol_alfa,
                   leg_angle_ouput[2].sol_alfa,
                   leg_angle_ouput[0].sol_beta,
                   leg_angle_ouput[1].sol_beta,
                   leg_angle_ouput[3].sol_beta,
                   leg_angle_ouput[2].sol_beta};

               setsmoothly(&robot, motors_target_pos, 0.1);
               robot.step(timeStep);
          }

// for (int i = 0; i < 5; i++)
// {
#ifdef MPC_TEST

          double t = robot.getTime();
          gaitPlanner.setGait(0.5, 0.5, "trot");
          gaitPlanner.calculateGait(t);

          // 获取IMU数据
          Vector3d Eulerangles = Map<Vector3d>(const_cast<double *>(iu->getRollPitchYaw())); // BUG: IMU零系不是通常理解的z轴为偏航
          Vector3d a_original = Map<Vector3d>(const_cast<double *>(acc->getValues()));
          Vector3d omega_original = Map<Vector3d>(const_cast<double *>(gyro->getValues()));

          // 编码器获取电机位置和速度
          Matrix<double, 4, 3> q_j_i;
          Matrix<double, 4, 3> q_dot_j_i;
          for (int i = 0; i < 4; i++)
          {
               q_j_i(i, 0) = ps[3 * i]->getValue();
               q_j_i(i, 1) = ps[3 * i + 1]->getValue();
               q_j_i(i, 2) = ps[3 * i + 2]->getValue();
               q_dot_j_i(i, 0) = motors[3 * i]->getVelocity();
               q_dot_j_i(i, 1) = motors[3 * i + 1]->getVelocity();
               q_dot_j_i(i, 2) = motors[3 * i + 2]->getVelocity();
          }

          stateEstimator.setIMUdata(Eulerangles, a_original, omega_original);
          stateEstimator.setEncoderdata(q_j_i, q_dot_j_i);
          stateEstimator.calculateAll();
          stateEstimator.kalman_update();

          Matrix<double, 18, 1> vecx_k = stateEstimator.getx_k();

          Matrix<double, 13, 1> vecx_0;
          vecx_0.block<3, 1>(0, 0) = Eulerangles;
          vecx_0.block<3, 1>(3, 0) = vecx_k.block<3, 1>(0, 0);
          vecx_0.block<3, 1>(6, 0) = omega_original;
          vecx_0.block<3, 1>(9, 0) = vecx_k.block<3, 1>(3, 0);
          vecx_0(12, 0) = -9.8;

          double v_x = 0.0, v_y = 0.0, omega_z = 0.0; // 每一周期的期望值
          // cout << "预期值: "
          //      << "v_x:" << v_x << " v_y:" << v_y << " omega_z:" << omega_z << endl;
          //  1. 将手柄传来的数据转化为质心轨迹

          Matrix<double, 13 * PREDICTION_HORIZON, 1> vecD;
          Matrix<double, PREDICTION_HORIZON, 1> vecpsi_k_d;

          trajectoryGenerator.setState(stateEstimator.getp_stend_O(gaitPlanner.gets_fai()));
          trajectoryGenerator.setDesired(v_x, v_y, omega_z);           // 设置期望值
          trajectoryGenerator.calculateAll(stateEstimator.getR_B_O()); // 计算质心轨迹
          trajectoryGenerator.getD(vecD);                              // 得到D向量
          trajectoryGenerator.getpsi_k_d(vecpsi_k_d);                  // 得到psi_k_d向量

          Matrix<double, 3 * PREDICTION_HORIZON, 4> r_k_ix;
          for (int i = 0; i < 4; i++)
          {
               if (gaitPlanner.gets_fai()(i) == 1)
               {
                    r_k_ix.col(i) = stateEstimator.getp_foot().row(i).transpose().replicate(1, PREDICTION_HORIZON);
               }
               else
               {
                    r_k_ix.col(i) = gaitPlanner.getp_swend_i_O(i).replicate(1, PREDICTION_HORIZON);
               }
          }

          mpcSolver.setD_and_x0(vecD, vecx_0); // 设置D向量
          mpcSolver.calculateAll(vecpsi_k_d, r_k_ix);
          mpcSolver.qpSolveU_first();

          // if (i == 4)
          // {
          //      // cout << "------------------x_0-----------------:" << endl
          //      // << vecx_0 << endl;
          //      // cout << "------------------x_k-----------------:" << endl
          //      //      << vecx_k << endl;
          //      // cout << "四足坐标: " << endl;
          //      // for (int i = 0; i < 4; i++)
          //      // {
          //      //      cout << "第" << i + 1 << "足: " << endl;
          //      //      cout << "x: " << vecx_k(6 + 3 * i) << " y: " << vecx_k(7 + 3 * i) << " z: " << vecx_k(8 + 3 * i) << endl;
          //      // }
          //      // cout << "------------------v_com-----------------:" << endl
          //      //      << vecx_k(3) << "  " << vecx_k(4) << "  " << vecx_k(5) << endl;
          //      // cout << "------------------D-----------------:" << endl
          //      //      << vecD << endl;

          //      // cout << "------------------x_0-----------------:" << endl;
          //      // cout << "位置:" << vecx_0.segment(3, 3).transpose() << endl;
          //      // cout << "速度:" << vecx_0.segment(9, 3).transpose() << endl;
          //      // cout << "角度: " << vecx_0.segment(0, 3).transpose() << endl;
          //      // cout << "角速度: " << vecx_0.segment(6, 3).transpose() << endl;

          //      // cout << "--------------------------------D--------------------------------:" << endl<<endl;
          for (int j = 0; j < PREDICTION_HORIZON; j++)
          {
               cout << "-----------------------------第 " << j + 1 << " 周期----------------------------------------" << endl;
               cout << setw(15) << "期望位置:" << std::setw(15) << vecD.segment(13 * j + 3, 3).transpose() << endl;
               cout << setw(15) << "期望速度:" << std::setw(15) << vecD.segment(13 * j + 9, 3).transpose() << endl;
               cout << setw(15) << "期望角度:" << std::setw(15) << vecD.segment(13 * j + 0, 3).transpose() << endl;
               cout << setw(15) << "期望角速度:" << std::setw(15) << vecD.segment(13 * j + 6, 3).transpose() << endl;
          }
          // }
          robot.step(TIME_STEP);
// }
#endif
     }
}

void reset_to_zero()
{
     cout << "reset_to_zero()" << endl;
     Robot *robot = new Robot();

     for (int i = 0; i < 12; i++)
     {
          Motor *motor = robot->getMotor(motor_names[i]);
          motor->setPosition(0);
          motor->getPositionSensor()->enable(TIME_STEP);
     }
     robot->step(TIME_STEP);

     delete robot;
}

int main()
{
     // test_controller();
     // cout << "before walk_bezier_forward()" << endl;
     // reset_to_zero();
     walk_bezier_forward();

     // spot_IMUtest();
     // MPCdemo();
     return 0;
}
